import os
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
from util import Parameter

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 标准化
])
transform_flip=transforms.Compose([
    transforms.RandomHorizontalFlip(p=Parameter.FLIP_SCALE),
    transform
])

class FaceDataset(Dataset):
    def __init__(self, root_dir, transform=transform_flip):
        """
        自定义数据集类
        :param root_dir: 包含图像文件的根目录
        :param transform: 应用于图像的转换
        """
        self.root_dir = root_dir
        self.transform = transform
        self.class_to_idx =Parameter.class_to_idx
        self.class_names = list(self.class_to_idx.keys())  # 获取类别名称
        self.image_paths = []  # 存储图像路径
        self.labels = []  # 存储图像标签

        # 遍历每个类别文件夹，获取图像路径和标签
        for cls_name in self.class_names:
            cls_dir = os.path.join(root_dir, cls_name)
            if not os.path.isdir(cls_dir):
                continue
            for img_name in os.listdir(cls_dir):
                img_path = os.path.join(cls_dir, img_name)
                if os.path.isfile(img_path):
                    self.image_paths.append(img_path)
                    self.labels.append(self.class_to_idx[cls_name])

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        label = self.labels[idx]
        image = Image.open(img_path)#.convert('RGB')  # 打开图像并转换为 RGB 格式

        if self.transform:
            image = self.transform(image)

        return image, label



# # 创建数据集实例
# dataset = CustomDataset(root_dir='./data', transform=transform)

# # 创建 DataLoader
# data_loader = DataLoader(dataset, batch_size=32, shuffle=True)

# # 测试 DataLoader
# for images, labels in data_loader:
#     print(f'Batch shape: {images.shape}, Labels: {labels}')
#     break